import cv2
import numpy as np
import threading
from queue import Queue
from paddleocr import PaddleOCR

class OCRVideoWrapper:
    def __init__(self, lang='ch'):
        self.ocr = PaddleOCR(use_angle_cls=True, lang=lang)
        self.frame_queue = Queue(maxsize=10)
        self.result_queue = Queue(maxsize=10)
        self.stop_flag = False

    def read_frames(self, cap):
        while not self.stop_flag and cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                self.stop_flag = True
                break
            self.frame_queue.put(frame)

    def process_frames(self):
        while not self.stop_flag or not self.frame_queue.empty():
            try:
                frame = self.frame_queue.get(timeout=1)
                result = self.ocr.ocr(frame, cls=True)
                self.result_queue.put((frame, result))
            except:
                continue

    def recognize_video(self, video_path, output_path=None, display=True):
        cap = cv2.VideoCapture(video_path)

        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        fps = cap.get(cv2.CAP_PROP_FPS)
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

        out = None
        if output_path:
            out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))

        # 启动线程
        read_thread = threading.Thread(target=self.read_frames, args=(cap,))
        process_thread = threading.Thread(target=self.process_frames)

        read_thread.start()
        process_thread.start()

        # 显示/保存结果
        while not self.stop_flag or not self.result_queue.empty():
            try:
                frame, result = self.result_queue.get(timeout=1)
                for line in result:
                    for word_info in line:
                        text = word_info[1][0]
                        points = word_info[0]
                        print(f"文本: {text}, 位置: {points}")
                        cv2.polylines(frame, [np.array(points, dtype=np.int32)], isClosed=True, color=(0, 255, 0), thickness=2)
                        cv2.putText(frame, text, (int(points[0][0]), int(points[0][1])),
                                    cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2)

                if display:
                    cv2.imshow("Video OCR", frame)
                    if cv2.waitKey(1) & 0xFF == ord('q'):
                        self.stop_flag = True
                        break

                if out:
                    out.write(frame)
            except:
                continue

        read_thread.join()
        process_thread.join()
        cap.release()
        if out:
            out.release()
        cv2.destroyAllWindows()
